// implement fork from user space

#include "lib.h"
#include <mmu.h>
#include <env.h>


/* ----------------- help functions ---------------- */

/* Overview:
 * 	Copy `len` bytes from `src` to `dst`.
 *
 * Pre-Condition:
 * 	`src` and `dst` can't be NULL. Also, the `src` area 
 * 	 shouldn't overlap the `dest`, otherwise the behavior of this 
 * 	 function is undefined.
 */
void user_bcopy(const void *src, void *dst, size_t len)
{
	void *max;

	//	writef("~~~~~~~~~~~~~~~~ src:%x dst:%x len:%x\n",(int)src,(int)dst,len);
	max = dst + len;

	// copy machine words while possible
	if (((int)src % 4 == 0) && ((int)dst % 4 == 0)) {
		while (dst + 3 < max) {
			*(int *)dst = *(int *)src;
			dst += 4;
			src += 4;
		}
	}

	// finish remaining 0-3 bytes
	while (dst < max) {
		*(char *)dst = *(char *)src;
		dst += 1;
		src += 1;
	}

	//for(;;);
}

/* Overview:
 * 	Sets the first n bytes of the block of memory 
 * pointed by `v` to zero.
 * 
 * Pre-Condition:
 * 	`v` must be valid.
 *
 * Post-Condition:
 * 	the content of the space(from `v` to `v`+ n) 
 * will be set to zero.
 */
void user_bzero(void *v, u_int64_t n)
{
	char *p;
	int m;

	p = v;
	m = n;

	while (--m >= 0) {
		*p++ = 0;
	}
}
/*--------------------------------------------------------------*/

/* Overview:
 * 	Custom page fault handler - if faulting page is copy-on-write,
 * map in our own private writable copy.
 * 
 * Pre-Condition:
 * 	`va` is the address which leads to a TLBS exception.
 *
 * Post-Condition:
 *  Launch a user_panic if `va` is not a copy-on-write page.
 * Otherwise, this handler should map a private writable copy of 
 * the faulting page at correct address.
 */
/*** exercise 4.13 ***/
void
pgfault(u_int64_t va)
{
	extern struct Env *envs;
        struct Env *thisenv;
        va = ROUNDDOWN(va, BY2PG);
        u_int64_t perm = 0x12345678;
	Pte *vpt2, *vpt1, *vpt0;
	u_int64_t tmpvar;
        u_int64_t tmp_addr = UTOP - 2 * BY2PG;
writef("fork.c:pgfault():\t va:%lx\n",va);
	/*u_int64_t id = syscall_getenvid();
	id = ENVX(syscall_getenvid());
	thisenv = envs+id;*/
	thisenv = &envs[ENVX(syscall_getenvid())];//&envs[i];
	vpt2 = thisenv->env_vpt2;
writef("fork.c: vpt2:0x%lx\n", vpt2);
	if ((vpt2[VPN2(va)] & PTE_V) == PTE_V) {
		vpt1 = PTE_TO_PADDR(vpt2[VPN2(va)]);
		if ((vpt1[VPN1(va)] & PTE_V) == PTE_V) {
			vpt0 = PTE_TO_PADDR(vpt1[VPN1(va)]);
			if ((vpt0[VPN0(va)] & PTE_V) == PTE_V) {
//				addr = (i << PT2SHIFT) + (j << PT1SHIFT) + (k << PT0SHIFT);
				perm = vpt0[VPN0(va)] & 0x3ff;
			}
		}
	}
//writef("Step1:perm = %lx\n", perm);
	if ((perm == 0x12345678) || ((perm & PTE_COW) == 0)) {
writef("Not COW!\n");
                user_panic("Error: page not copy on write!");
	}
/*	if ((((Pte *)(*vpt))[VPN(va)] & PTE_COW) == 0) {
                user_panic("Error: page not copy on write!");
        }*/
    //map the new page at a temporary place
	tmpvar = syscall_mem_alloc(0, tmp_addr, PTE_U | PTE_V | PTE_R | PTE_W | (perm & (~PTE_COW)));
//writef("Step2\n");
        if (tmpvar != 0) {
                user_panic("Error: map to tmp addr error!");
        }
        //copy the content
        user_bcopy(va, tmp_addr, BY2PG);
    //map the page on the appropriate place
        if (syscall_mem_map(0, tmp_addr, 0, va, PTE_U | PTE_V | PTE_R | PTE_W | (perm & (~PTE_COW))) != 0) {
                user_panic("Error: map to appropriate place error!");
        }
        if (syscall_mem_unmap(0, tmp_addr) != 0) {
                user_panic("Error: unmap from tmp addr error!");
        }
    //unmap the temporary place
writef("pgfault end\n");
}

/* Overview:
 * 	Map our virtual page `pn` (address pn*BY2PG) into the target `envid`
 * at the same virtual address. 
 *
 * Post-Condition:
 *  if the page is writable or copy-on-write, the new mapping must be 
 * created copy on write and then our mapping must be marked 
 * copy on write as well. In another word, both of the new mapping and
 * our mapping should be copy-on-write if the page is writable or 
 * copy-on-write.
 * 
 * Hint:
 * 	PTE_LIBRARY indicates that the page is shared between processes.
 * A page with PTE_LIBRARY may have PTE_R at the same time. You
 * should process it correctly.
 */
/*** exercise 4.10 ***/
void
duppage(u_int64_t envid, u_int64_t addr, u_int64_t perm)
{
        //u_int64_t addr;
        //u_int64_t perm;
        //addr = pn * BY2PG;
        //perm = ((Pte *)(*vpt))[pn] & 0xfff;
//writef("duppage:%x\n", addr);
	int r;
	if ((perm & PTE_V) == 0) {
		user_panic("duppage while PTE_V not valid! Addr:%lx", addr);
	}
        if ((perm & PTE_LIBRARY) != 0) {
                perm = perm | PTE_W | PTE_U;
        } else if (((perm & PTE_W) != 0) || ((perm & PTE_COW) != 0)) {
                perm = perm | PTE_COW | PTE_U;
		perm = perm | PTE_W;
		perm = perm ^ PTE_W;
        }
	
//writef("map childenv\n");
        if (syscall_mem_map(0, addr, envid, addr, perm) < 0) {
                user_panic("duppage err at son!");
        }
	r = syscall_mem_map(0, addr, 0, addr, perm);
//writef("map self\n");
        if (r < 0) {
                user_panic("duppage err at father!");
        }
	writef("duppage success at %lx, perm %lx\n", addr, perm);
        //      user_panic("duppage not implemented");
}

/* Overview:
 * 	User-level fork. Create a child and then copy our address space
 * and page fault handler setup to the child.
 *
 * Hint: use vpd, vpt, and duppage.
 * Hint: remember to fix "env" in the child process!
 * Note: `set_pgfault_handler`(user/pgfault.c) is different from 
 *       `syscall_set_pgfault_handler`. 
 */
/*** exercise 4.9 4.15***/
extern void __asm_pgfault_handler(void);
int
fork(void)
{
        // Your code here.
        u_int64_t newenvid;
        extern struct Env *envs;
        extern struct Env *env;
        u_int64_t i, j, k, addr=0, perm;
	Pte *vpt2, *vpt_1, *vpt0;

        //The parent installs pgfault using set_pgfault_handler
        set_pgfault_handler(pgfault);
        //alloc a new alloc
        newenvid = syscall_env_alloc();
writef("Newenv got\n");
        if (newenvid == 0) {
                env = &envs[ENVX(syscall_getenvid())];//&envs[i];
writef("son:envaddr = %lx\n");
                return 0;
        }
/*	for (i = 0; i < UTOP; i = i + PDMAP) {
		if(((Pde *)(*vpd))[i >> PDSHIFT] != 0) {
			for(j = 0; (j < PDMAP) && ((i + j) < USTACKTOP); j += BY2PG) {
				if (((Pte *)(*vpt))[VPN(i + j)] != 0) {
					duppage(newenvid, VPN(i + j));
				}
			}
		}
	}*/
	vpt2 = env->env_vpt2;
writef("father, PTSIZE=%lx, envaddr=%lx\n", PTSIZE, env);

	for (i = 0; i <= VPN2(UTOP); i++) {
		if ((vpt2[i] & PTE_V) == PTE_V) {
			vpt_1 = (Pte *)PTE_TO_PADDR(vpt2[i]);
			for (j = 0; j < PTSIZE; j++) {
				if ((vpt_1[j] & PTE_V) == PTE_V) {
					vpt0 = PTE_TO_PADDR(vpt_1[j]);
					for (k = 0; k < PTSIZE; k++) {
						if ((vpt0[k] & PTE_V) == PTE_V) {
							addr = (i << PT2SHIFT) | (j << PT1SHIFT) | (k << PT0SHIFT);
							perm = vpt0[k] & 0x3ff;
							if (addr >= UTOP) {
								break;
							}
							duppage(newenvid, addr, perm);
						}
					}
				}
			}
		}
	}
writef("duppage success!\n");
	if (syscall_mem_alloc(newenvid, UXSTACKTOP - BY2PG, PTE_V | PTE_R | PTE_LIBRARY | PTE_W | PTE_U) != 0) {
		user_panic("User exception stack error!");
	}
writef("Exc stack set success!\n");
	if (syscall_set_pgfault_handler(newenvid, __asm_pgfault_handler, UXSTACKTOP) != 0) {
                        user_panic("Syscall_set_pgfault_handler error!");
	}
	if (syscall_set_env_status(newenvid, ENV_RUNNABLE) != 0) {
		user_panic("syscall set env status error!");
	}

        return newenvid;
}

// Challenge!
int
sfork(void)
{
	user_panic("sfork not implemented");
	return -E_INVAL;
}
